package com.zx.gateway.filter;

import com.zx.core.utils.AES128CBCNoPadding;
import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.cloud.context.config.annotation.RefreshScope;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.server.reactive.ServerHttpRequest;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.util.StringUtils;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.charset.StandardCharsets;

/**
 * <p>
 * description: 自定义过滤器对请求和响应进行加密解密 <br>
 * create: 2025-01-08 15:39
 * </p>
 *
 * @author zhou  xun
 */
@Slf4j
@RefreshScope //支持Nacos的动态刷新功能。
@Component
public class EncryptDecryptFilter implements GlobalFilter, Ordered {

    @Value("${security.encrypt.enabled}")
    private Boolean encryptEnabled;

    @Value("${security.encrypt.key}")
    private String encryptKey;

    @Value("${security.encrypt.iv}")
    private String encryptIv;

    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {
        // 如果未启用加密功能，则直接跳过
        if (!encryptEnabled) {
            return chain.filter(exchange);
        }

        // 检查加密所需的配置是否完整
        if (!StringUtils.hasText(encryptKey) || !StringUtils.hasText(encryptIv)) {
            // 配置不完整，记录警告并跳过加密处理
            return chain.filter(exchange);
        }

        ServerHttpRequest originalRequest = exchange.getRequest();
        DataBufferFactory bufferFactory = exchange.getResponse().bufferFactory();

        return DataBufferUtils.join(originalRequest.getBody())
                .flatMap(dataBuffer -> {
                    if (dataBuffer.readableByteCount() == 0) {
                        return handleResponse(chain, exchange, originalRequest);
                    }

                    byte[] content = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(content);
                    DataBufferUtils.release(dataBuffer);

                    String encryptedData = new String(content, StandardCharsets.UTF_8);
                    String decryptedData;
                    try {
                        // 记录解密前的加密数据
                        log.info("Decrypting request body - Encrypted: {}", encryptedData);
                        decryptedData = AES128CBCNoPadding.decrypt(encryptKey, encryptIv, encryptedData);
                        // 记录解密后的明文数据
                        log.info("Decrypting request body - Decrypted: {}", decryptedData);
                    } catch (Exception e) {
                        // 日志记录错误信息
                        log.error("Failed to decrypt request body: {}", encryptedData, e);
                        return Mono.error(new RuntimeException("Decryption failed", e));
                    }

                    // 创建带有解密后内容的新请求
                    Flux<DataBuffer> modifiedBody = Flux.just(
                            bufferFactory.wrap(decryptedData.getBytes(StandardCharsets.UTF_8))
                    );

                    // 构建新的 ServerHttpRequestDecorator 对象
                    ServerHttpRequest mutatedRequest = new ServerHttpRequestDecorator(originalRequest) {
                        @Override
                        public Flux<DataBuffer> getBody() {
                            return modifiedBody;
                        }
                    };

                    // 使用新的请求对象继续过滤链，并处理响应加密
                    return handleResponse(chain, exchange.mutate().request(mutatedRequest).build(), mutatedRequest);
                });
    }

    private Mono<Void> handleResponse(GatewayFilterChain chain, ServerWebExchange exchange, ServerHttpRequest request) {
        ServerHttpResponse originalResponse = exchange.getResponse();
        DataBufferFactory bufferFactory = originalResponse.bufferFactory();

        ServerHttpResponseDecorator decoratedResponse = new ServerHttpResponseDecorator(originalResponse) {
            @Override
            public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
                // 如果未启用加密功能，则直接返回原始响应
                if (encryptEnabled) {
                    return super.writeWith(body);
                }

                // 检查加密所需的配置是否完整
                if (!StringUtils.hasText(encryptKey) || !StringUtils.hasText(encryptIv)) {
                    // 配置不完整，记录警告并返回原始响应
                    log.warn("加密配置不完整，返回原始响应数据");
                    return super.writeWith(body);
                }

                if (body instanceof Flux) {
                    Flux<? extends DataBuffer> flux = (Flux<? extends DataBuffer>) body;
                    return super.writeWith(flux.buffer().map(dataBuffers -> {
                        // 合并所有数据
                        byte[] combinedBytes = new byte[dataBuffers.stream().mapToInt(DataBuffer::readableByteCount).sum()];
                        int offset = 0;
                        for (DataBuffer dataBuffer : dataBuffers) {
                            int length = dataBuffer.readableByteCount();
                            dataBuffer.read(combinedBytes, offset, length);
                            offset += length;
                        }

                        String responseData = new String(combinedBytes, StandardCharsets.UTF_8);
                        log.debug("Original response data: {}", responseData);

                        try {
                            // 加密响应数据
                            String encryptedResponse = AES128CBCNoPadding.encrypt(encryptKey, encryptIv, responseData);
                            log.debug("Encrypted response data: {}", encryptedResponse);

                            // 返回加密后的数据
                            return bufferFactory.wrap(encryptedResponse.getBytes(StandardCharsets.UTF_8));
                        } catch (Exception e) {
                            log.error("Failed to encrypt response body", e);
                            throw new RuntimeException("Encryption failed", e);
                        }
                    }));
                }
                return super.writeWith(body);
            }
        };

        // 继续过滤器链
        return chain.filter(exchange.mutate().request(request).response(decoratedResponse).build());
    }

    @Override
    public int getOrder() {
        return -1;
    }
}